package jtg.graphics;

import soot.*;
import soot.jimple.*;
import soot.jimple.internal.ImmediateBox;
import soot.options.Options;
import soot.toolkits.graph.ClassicCompleteUnitGraph;
import soot.toolkits.graph.UnitGraph;
import soot.util.Chain;

import java.io.File;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public class SootCFG {


    public static UnitGraph getMethodCFG(String sourceDirectory,String clsName,String methodName){
        Body body = getMethodBody(sourceDirectory,clsName,methodName);
        UnitGraph ug = new ClassicCompleteUnitGraph(body);
        return ug;
    }

    public static Body getMethodBody(String sourceDirectory,String clsName,String methodName){
        G.reset();
        Options.v().set_prepend_classpath(true);
        Options.v().set_allow_phantom_refs(true);
        Options.v().set_soot_classpath(sourceDirectory);
        SootClass sc = Scene.v().loadClassAndSupport(clsName);
        sc.setApplicationClass();
        Scene.v().loadNecessaryClasses();
        SootClass mainClass = Scene.v().getSootClass(clsName);
        SootMethod sm = mainClass.getMethodByName(methodName);
        Body body = sm.retrieveActiveBody();
        return body;
    }

    public static void addPrefixToValueNames(Body body, String prefix) {
        for (Local local : body.getLocals()) {
            local.setName(prefix + local.getName());
        }
    }

    public static void addArgs(Body body, List<Value> args) {
        int cnt = 0;
        for (Iterator<Unit> it = body.getUnits().snapshotIterator(); it.hasNext(); ) {
            Unit unit = it.next();
            if (unit instanceof IdentityStmt) {
                IdentityStmt identityStmt = (IdentityStmt) unit;
                Value leftValue = identityStmt.getLeftOp();
                Value rightValue = identityStmt.getRightOp();
                if (rightValue instanceof ParameterRef) {
                    // 增加新的AssignStmt
                    AssignStmt assignStmt = Jimple.v().newAssignStmt(leftValue, args.get(cnt));
                    body.getUnits().insertBefore(assignStmt, unit);
                    // 删除原先的IdentityStmt
                    body.getUnits().remove(unit);
                    cnt++;
                }
            }
        }
    }

    public static void inlineMethodBody(Body body) {
        Iterator<Unit> iterator = body.getUnits().snapshotIterator();
        while (iterator.hasNext()) {
            Unit unit = iterator.next();
            if (unit instanceof AssignStmt){
                AssignStmt assignStmt = (AssignStmt) unit;
                Value leftValue = assignStmt.getLeftOp();
                Value rightValue = assignStmt.getRightOp();
                if(rightValue instanceof InvokeExpr) {
                    InvokeExpr invokeExpr = (InvokeExpr) rightValue;
                    SootMethod sm = invokeExpr.getMethod();
                    if(sm.isConcrete() && !sm.isJavaLibraryMethod()){
                        List<Value> callArgs = invokeExpr.getArgs();
                        Body calleeBody = (Body) sm.retrieveActiveBody().clone();
                        // 替换 parameter 为实际参数
                        addArgs(calleeBody, callArgs);
                        // 替换 calleeBody 中的变量名, 防止变量名冲突
                        addPrefixToValueNames(calleeBody,sm.getName() + "_");
                        Iterator<Unit> calleeIterator = calleeBody.getUnits().snapshotIterator();
                        while (calleeIterator.hasNext()) {
                            Unit calleeUnit = calleeIterator.next();
                            if (calleeUnit instanceof ReturnStmt) {
                                ReturnStmt returnStmt = (ReturnStmt) calleeUnit;
                                Value returnValue = returnStmt.getOp();
                                AssignStmt newAssignStmt = Jimple.v().newAssignStmt(leftValue,returnValue);
                                body.getUnits().insertAfter(newAssignStmt,unit);
                                calleeBody.getUnits().remove(calleeUnit);
                            }
                        }
                        body.getUnits().insertAfter(calleeBody.getUnits(),unit);
                        body.getUnits().remove(unit);
                    }
                }
            }
        }
    }

}
